# Copyright (c) 2015 Presto Labs Pte. Ltd.
# Author: jaewon

import struct


class RecordWriter(object):
  def __init__(self, filename):
    self._filename = filename
    self._file = open(filename, 'wb')

  def __enter__(self):
    return self

  def __exit__(self, type, value, traceback):
    self._flle.close()

  def open(self):
    return True

  def close(self):
    self._file.close()

  def flush(self):
    self._file.flush()

  def write_record(self, buf):
    if type(buf) is bytes:
      pass
    elif type(buf) is str:
      buf = str.encode(buf)
    else:
      return False
    size = len(buf)
    size_buf = struct.pack('<i', size)
    self._file.write(size_buf)
    self._file.write(buf)
    return True


class RecordReader(object):
  def __init__(self, filename):
    self._filename = filename
    self._file = open(filename, 'rb')
    self._last_pos = 0
    self._record_offset = 0

  def __enter__(self):
    return self

  def __exit__(self, type, value, traceback):
    self._flle.close()

  def open(self):
    return True

  def close(self):
    self._file.close()

  def reopen(self):
    if self._file:
      self._file.close()
    self._file = open(filename, 'rb')
    self._file.seek(self._last_pos)

  def read_record(self):
    buf = self._file.read(4)
    if len(buf) < 4:
      self._file.seek(self._last_pos)
      return None

    record_size = struct.unpack('<i', buf)[0]
    buf = self._file.read(record_size)
    if len(buf) < record_size:
      self._file.seek(self._last_pos)
      return None

    self._last_pos = self._file.tell()
    self._record_offset += 1
    return buf

  def read_records(self):
    while True:
      record = self.read_record()
      if record:
        yield record
      else:
        break

  @property
  def last_record_offset(self):
    return self._record_offset - 1
